import yaml
import os
import time
from configs.config import global_var
import math
import json


def load_cfg(path):
    with open(path, 'r') as f:
        return yaml.load(f, Loader=yaml.FullLoader)


def mkdir(path):
    folder = os.path.exists(path)

    if not folder:  # 判断是否存在文件夹如果不存在则创建为文件夹
        os.makedirs(path)  # makedirs 创建文件时如果路径不存在会创建这个路径
        print("---  new folder...  ---")
        print("---  OK  ---")
    else:
        print("---  There is this folder!  ---")


def new_fold(mode, filename_list, name=None):
    # 首先生成新的文件夹用来保存测试数据
    if name is not None:
        prefix = name
    else:
        prefix = time.strftime("%Y-%m-%d-%H-%M", time.localtime(time.time()))
    global_var.saving_address = global_var.saving_address + mode + '/' + prefix + '/'
    for filename in filename_list:
        file_path = global_var.saving_address + filename
        mkdir(file_path)


def quart_normalization(q1, q2, q3, q4):
    # 判断四元素的合理性，即平方和为1
    para_sum = q1 ** 2 + q2 ** 2 + q3 ** 2 + q4 ** 2
    q1 = (q1 ** 2) / para_sum
    q2 = (q2 ** 2) / para_sum
    q3 = (q3 ** 2) / para_sum
    q4 = (q4 ** 2) / para_sum
    return q1, q2, q3, q4


def quart_to_rpy(x, y, z, w):
    roll = math.atan2(2 * (w * x + y * z), 1 - 2 * (x * x + y * y))
    pitch = math.asin(2 * (w * y - x * z))
    yaw = math.atan2(2 * (w * z + x * y), 1 - 2 * (z * z + y * y))
    roll = round(roll / math.pi * 180, 3)
    pitch = round(pitch / math.pi * 180, 3)
    yaw = round(yaw / math.pi * 180, 3)
    return roll, pitch, yaw


# 用于统计网络的可训练参数个数
def get_parameter_number(model):
    total_num = sum(p.numel() for p in model.parameters())
    trainable_num = sum(p.numel() for p in model.parameters() if p.requires_grad)
    return {'Total': total_num, 'Trainable': trainable_num}


def batch_file_name_change(path='C:/Users/adminTKJ/Desktop/RLIR_sumup/data/CT/均匀采样/train/标准正位/DRR'):
    # 获取该目录下所有文件，存入列表中
    fileList = os.listdir(path)
    n = 0
    for i in fileList:
        # 设置旧文件名（就是路径+文件名）
        oldname = path + os.sep + fileList[n]  # os.sep添加系统分隔符

        # 设置新文件名
        newname = path + os.sep + i[:-4]

        os.rename(oldname, newname)  # 用os模块中的rename方法对文件改名
        print(oldname, '======>', newname)

        n += 1


#
def get_list_from_dict(_dict, _id):
    _list = []
    for _key, _value in _dict.items():
        try:
            _list.append(round(float(_value[_id]), 2))
        except IndexError:
            # 噪声的格式不一样，需要特殊处理
            _noise = _value[0] ** 2 + _value[1] ** 2
            _list.append(round(_noise, 2))

    return _list


def amend_json():
    json_path = 'C:/Users/adminTKJ/Desktop/RLIR_sumup/data/CT/均匀采样/train/标准侧位/label.json'
    # 读取json文件并修改键
    with open(json_path, 'r') as raw_file:
        new_data = json.load(raw_file)
        for key, value in new_data.items():
            print('修改前：')
            print(new_data[key])
            value["angle_y"] = value.pop("angel_y")
            value["angle_z"] = value.pop("angel_z")
            print(value)
            print('修改后：')
            print(new_data[key])
    with open(json_path, 'w') as new_file:
        json.dump(new_data, new_file)


if __name__ == "__main__":
    batch_file_name_change()
